#include <stdio.h>
#include <stdlib.h>

#define EXTRA_INFO_SZ 4
#define _EXTRA_INFO_ALIGN 4
#define EXTRA_INFO_ALIGN (_EXTRA_INFO_ALIGN-1)
#define EXTRA_INFO_TOTAL (EXTRA_INFO_SZ + _EXTRA_INFO_ALIGN)

#define alignPtr(ptr,align) (void *) ( ((size_t)(ptr) + (align)) & ~(size_t)(align) )

typedef struct
{
  char *base;
  char *mem;
  unsigned int totalSz;
  unsigned int usedSz;
} stackMem;

typedef struct
{
  stackMem *stacks;
  unsigned int nStacks;
  unsigned int align;
} totalMem;

size_t totalSz(totalMem *tm)
{
  int i;
  size_t totalSz = 0;

  for(i=0; i<tm->nStacks; i++) {
    totalSz += tm->stacks[i].totalSz;
  }

  return totalSz;
}

int addStackMem(totalMem *tm, int minSz)
{
  int i;
  size_t newStackSz = totalSz(tm);

  if(newStackSz<minSz) newStackSz = minSz;

  stackMem *stacks = realloc(tm->stacks,sizeof(stackMem)*(tm->nStacks+1));
  if(!stacks) return 0;

  tm->stacks = stacks;
  stackMem *sm = stacks+tm->nStacks;

  sm->base = (char *)malloc(newStackSz+tm->align);
  sm->mem = alignPtr(sm->base,tm->align);
  sm->totalSz = (newStackSz+tm->align)-(sm->mem-sm->base);
  sm->usedSz = 0;

  tm->nStacks++;

  return 1;
}

totalMem *initMem(int initSz, int align)
{
  totalMem *tm = (totalMem *)malloc(sizeof(totalMem));

  if(tm) {
    tm->stacks = 0;
    tm->nStacks = 0;
    tm->align = align - 1;

    if(!addStackMem(tm, initSz)) {
      free(tm);
      tm = 0;
    }
  }

  return tm;
}

void *alloc(totalMem *tm, int sz)
{
  stackMem *sm = &tm->stacks[tm->nStacks-1];

  int freeMem = sm->totalSz - sm->usedSz;
  if(freeMem < (sz+EXTRA_INFO_TOTAL)) {
    if(!addStackMem(tm,sz+EXTRA_INFO_TOTAL)) return 0;
    sm = &tm->stacks[tm->nStacks-1];
  }

  char *result = sm->mem+sm->usedSz;

  char *extraInfo = result + sz;
  char *extraAligned = (char *)alignPtr(extraInfo,EXTRA_INFO_ALIGN);

  *((unsigned int *)extraAligned) = sm->usedSz;

  sm->usedSz += (extraAligned+EXTRA_INFO_SZ)-result;

  return result;
}

void *allocAlign(totalMem *tm, int sz, int align)
{
  align-=1;
  void *resultTmp = alloc(tm,sz+align);

  return alignPtr(resultTmp,align);
}

void dealloc(totalMem *tm)
{
  stackMem *sm = &tm->stacks[tm->nStacks-1];

  unsigned int *lastExtraInfo = (unsigned int *) (sm->mem + sm->usedSz - EXTRA_INFO_SZ);
  sm->usedSz = *lastExtraInfo;

  if(sm->usedSz==0) {
    free(sm->base);
    tm->nStacks--;
  }
}

//#define MS 1
#define MS 1024
//#define MS 1024*1024

void main(void)
{
  totalMem *mem = initMem(256,32);

  alloc(mem,128*MS);     //+
  alloc(mem,128*MS);     //++

  dealloc(mem);          //+

  alloc(mem,120*MS);     //++
  alloc(mem,8*MS);       //+++
  alloc(mem,1024*MS);    //++++

  dealloc(mem);          //+++

  alloc(mem,1023*MS);    //++++
  allocAlign(mem,1*MS,4);//+++++
  allocAlign(mem,1*MS,4);//++++++

  dealloc(mem);          //+++++
  dealloc(mem);          //++++
  dealloc(mem);          //+++
  dealloc(mem);          //++
  dealloc(mem);          //+
  dealloc(mem);          //
}

